open Core.Std
open Ast
open Semantics
open Printf

type module_type =
  | Whitebox
  | Blackbox of (string * int) list (* parameter list *)

type driver =
  | Driver_var of endpoint
  | Driver_const of int

type wire = driver * endpoint * int
type datapath = driver * endpoint * int * (Int.Set.t option) (* src, dst, width, label *)
(* datapath of width 1 *)
type datapath1 = driver * endpoint * (Int.Set.t option)

let string_of_path (p : path) : string =
  String.concat p ~sep:"__"

let string_of_endpoint (ep : endpoint) : string =
  let path, bit = ep in
  sprintf "%s__%d" (string_of_path path) bit

let expand_endpoint (path, lo) width =
  List.range 0 width
  |> List.map ~f:(fun i -> path, lo + i)

let expand_datapath dp =
  let src_drv, dst_ep, width, label = dp in
  let dst_path, dst_lo = dst_ep in
  match src_drv with
  | Driver_var src_ep ->
      let src_path, src_lo = src_ep in
      List.map (List.range 0 width) ~f:begin fun bit ->
        let src_ep' = (src_path, src_lo + bit) in
        let dst_ep' = (dst_path, dst_lo + bit) in
        Driver_var src_ep', dst_ep', label
      end
  | Driver_const src_lit ->
      List.map (List.range 0 width) ~f:begin fun bit ->
        let dst_ep' = (dst_path, dst_lo + bit) in
        Driver_const (src_lit lsr bit land 1), dst_ep', label
      end

let expand_datapaths datapaths =
  List.map datapaths ~f:expand_datapath
  |> List.concat

type node = {
  component_array : (string * component) array;
  lookup : string -> component;
}
and component =
  | Port of Vlog.port
  | Submodule of node

let rec map_ports (f : path * Vlog.port -> 'a) (node : node) : 'a list =
  node.component_array
  |> Array.to_list
  |> List.map ~f:begin fun (name, comp) ->
      match comp with
      | Port port -> [f ([name], port)]
      | Submodule node' ->
          map_ports (fun (path, port) -> f (name :: path, port)) node'
  end
  |> List.concat

type stage = {
  st_wires : (string * int) list; (* name, width *)
  st_instances : (string * Vlog.vlog_instance) list;
  st_assigns : Vlog.assign list;
}

let rec follow_path (node : node) (p : path) : Vlog.port =
  match p with
  | [] -> assert false
  | name :: p' ->
      let comp =
        try node.lookup name
        with Not_found ->
          failwithf "no such component: %s" (String.concat p ~sep:".") ()
      in
      match comp with
      | Port port -> port
      | Submodule node' -> follow_path node' p'

let subnode (node : node) (i : int) : node =
  match node.component_array.(i) with
  | _, Submodule node' -> node'
  | _ -> assert false

let convert_component_alist (alist : (string * component) list) : node =
  let component_array = Array.of_list alist in
  let lookup name =
    let name_map =
      alist
      |> List.mapi ~f:(fun index (name, _) -> (name, index))
      |> String.Map.of_alist_exn
    in
    snd component_array.(String.Map.find_exn name_map name)
  in
  { component_array; lookup }

let fail_with_parsing_error filename lexbuf =
  let curr = lexbuf.Lexing.lex_curr_p in
  let line = curr.Lexing.pos_lnum in
  let col = curr.Lexing.pos_cnum - curr.Lexing.pos_bol in
  failwithf "%s:%d:%d" filename line col ()

type blackbox = {
  class_name : string;
  ports : (string * Vlog.port) list;
  has_clock : bool;
}

let read_blackbox (search_dirs : string list) (class_name : string) : blackbox =
  let path =
    List.find_map search_dirs ~f:begin fun search_dir ->
      let path =sprintf "%s/bb/%s" search_dir class_name in
      if Sys.is_file_exn path
      then Some path
      else None
    end
    |> function
      | Some path -> path
      | None ->
          failwithf "module %s not found in following directories: %s"
            class_name (String.concat search_dirs ~sep:"; ") ()
  in
  let in_chan = open_in path in
  let lexbuf = Lexing.from_channel in_chan in
  let has_clock, ports =
    try Parser.bb_spec Lexer.read lexbuf
    with Parser.Error -> fail_with_parsing_error path lexbuf
  in
  { class_name; ports; has_clock }

let convert_components
  (module_search_dirs : string list)
  (port_alist : (string * Vlog.port) list)
  (blackboxes : ast_blackbox list)
  : (string * component) list =
  let blackbox_nodes =
    blackboxes
    |> List.map ~f:
      begin fun (_, (class_name, _)) ->
        let bbox = read_blackbox module_search_dirs class_name in
        List.map bbox.ports ~f:begin fun (port_name, port) ->
          (port_name, Port port)
        end
        |> convert_component_alist
      end
  in
  let port_component_alist : (string * component) list =
    port_alist
    |> List.map ~f:(fun (name, port) -> (name, Port port))
  in
  let submodule_component_alist =
    List.zip_exn blackboxes blackbox_nodes
    |> List.map ~f:(fun ((name, _), node) -> (name, Submodule node))
  in
  port_component_alist @ submodule_component_alist

exception Invalid_path

let resolve_endpoint (node : node) (ep : ast_endpoint) : endpoint * int =
  let rec resolve_path (node : node) (p : string list) : Vlog.port =
    match p with
    | [] -> raise Invalid_path (* path does not denote a port *)
    | name :: p' ->
        let comp =
          try node.lookup name
          with Not_found ->
            failwithf "no such component: %s"
              (String.concat ep.ep_path ~sep:".") ()
        in
        begin match comp with
        | Port port -> port
        | Submodule node' -> resolve_path node' p'
        end
  in
  let port_width, _ =
    try resolve_path node ep.ep_path
    with Invalid_path ->
      failwithf "Path does not denote a port: %s"
        (String.concat ep.ep_path ~sep:".") ()
  in
  let lo, hi = Option.value ep.ep_range ~default:(0, port_width) in
  if hi > port_width then failwithf "vector index out of bounds" ();
  ((ep.ep_path, lo), hi-lo)

let resolve_wire (node : node) (ast_wire : ast_wire) : wire =
  let src_ast_drv, dst_ast_ep = ast_wire in
  let dst_ep, dst_width = resolve_endpoint node dst_ast_ep in
  match src_ast_drv with
  | Ast_driver_var src_ast_ep ->
      let src_ep, src_width = resolve_endpoint node src_ast_ep in
      if src_width <> dst_width
      then begin
        let src_ast_ep_s = string_of_ast_endpoint src_ast_ep in
        let dst_ast_ep_s = string_of_ast_endpoint dst_ast_ep in
        failwithf "endpoint width mismatch: %s (%d) and %s (%d)"
          src_ast_ep_s src_width dst_ast_ep_s dst_width ()
      end;
      Driver_var src_ep, dst_ep, src_width
  | Ast_driver_const (src_width, src_value) ->
      (* TODO check width *)
      Driver_const src_value, dst_ep, src_width

let resolve_labels (lookup_label : string -> int) (labels : string list) : Int.Set.t =
  Int.Set.of_list (List.map labels ~f:lookup_label)

let convert_wire (node : node) (src_ast_drv, dst_ast_ep) : Vlog.assign =
  let (assign_rhs : Vlog.expr), (src_width : int) =
    match src_ast_drv with
    | Ast_driver_var src_ast_ep ->
        let src_ep, src_width = resolve_endpoint node src_ast_ep in
        let src_path, src_lo = src_ep in
        Vlog.Sym_part begin
          string_of_path src_path, (src_lo, src_lo + src_width)
        end,
        src_width
    | Ast_driver_const (width, value) ->
        Vlog.Literal (width, value), width
  in
  let dst_ep, dst_width = resolve_endpoint node dst_ast_ep in
  (* check width *)
  if dst_width <> src_width
  then begin
    let src_s =
      match src_ast_drv with
      | Ast_driver_var src_ast_ep -> string_of_path src_ast_ep.ep_path
      | Ast_driver_const _ -> "<const>"
    in
    let dst_s = string_of_path (fst dst_ep) in
    failwithf "endpoint width mismatch: %s (%d) and %s (%d)"
      src_s src_width dst_s dst_width ()
  end;
  let dst_path, dst_lo = dst_ep in
  let assign_lhs =
    string_of_path dst_path, (dst_lo, dst_lo + src_width)
  in
  assign_lhs, assign_rhs

exception Invalid_datapath

let resolve_datapath
  (node : node)
  (lookup_label : string -> int)
  (dp : ast_datapath)
  : datapath =
  let src_ast_drv, dst_ast_ep, ast_labels_opt = dp in
  let src_drv, dst_ep, width = resolve_wire node (src_ast_drv, dst_ast_ep) in
  let label_set_opt =
    match ast_labels_opt with
    | Some ast_labels -> Some (resolve_labels lookup_label ast_labels)
    | None -> None
  in
  src_drv, dst_ep, width, label_set_opt

let vlog_of_endpoint (ep : endpoint) : Vlog.expr =
  let path, bit = ep in
  Vlog.Sym_part (string_of_path path, (bit, bit+1))

let main (destdir : string) (search_paths : string list) () : unit =
  let lexbuf = Lexing.from_channel stdin in
  let top_ast =
    try
      Parser.top Lexer.read lexbuf
    with _ -> fail_with_parsing_error "<stdin>" lexbuf
  in
  let inst_num = List.length top_ast.top_encodings in
  let stage_names = top_ast.top_pipeline in
  let stage_name_arr = Array.of_list stage_names in
  let index_of_stage =
    stage_name_arr
    |> Array.to_list
    |> List.mapi ~f:(fun i name -> name, i)
    |> String.Map.of_alist_exn
    |> String.Map.find_exn
  in
  let first_stage_name = List.hd_exn stage_names in
  (* synthesize port "inst_class"on every stage *)
  List.iter top_ast.top_stages ~f:begin fun (_, stage_ast) ->
    let inst_class_port = ("inst_class", (inst_num, Vlog.Output)) in
    stage_ast.stage_ports <- inst_class_port :: stage_ast.stage_ports
  end;
  let sorted_ast_stage_alist =
    List.map top_ast.top_pipeline ~f:begin fun stage_name ->
      stage_name,
      List.Assoc.find_exn top_ast.top_stages stage_name ~equal:String.equal
    end
  in
  let stage_node_arr =
    List.map sorted_ast_stage_alist ~f:begin fun (_, st) ->
      convert_components search_paths st.stage_ports st.stage_modules
      |> convert_component_alist
    end
    |> Array.of_list
  in
  let stage_component_alist =
    List.mapi top_ast.top_pipeline ~f:begin fun i stage_name ->
      stage_name, Submodule stage_node_arr.(i)
    end
  in
  let top_blackboxes : (string * blackbox) list =
    List.map top_ast.top_modules ~f:begin fun (instance_name, (class_name, _)) ->
      instance_name, read_blackbox search_paths class_name
    end
  in
  let top_node =
    let port_component_alist =
      List.map top_ast.top_ports ~f:(fun (name, port) -> name, Port port)
    in
    let blackbox_component_alist =
      List.map top_blackboxes ~f:begin fun (bbox_name, bbox) ->
        let bbox_port_component_alist =
          List.map bbox.ports ~f:(fun (name, port) -> name, Port port)
        in
        bbox_name, Submodule (convert_component_alist bbox_port_component_alist)
      end
    in
    [ stage_component_alist; blackbox_component_alist; port_component_alist ]
    |> List.concat
    |> convert_component_alist
  in
  let lookup_label =
    top_ast.top_encodings
    |> List.mapi ~f:(fun i label -> label, i)
    |> String.Map.of_alist_exn
    |> begin fun map label ->
      try String.Map.find_exn map label
      with Not_found -> failwithf "No such instruction: %s" label ()
    end
  in
  let stages =
    List.mapi sorted_ast_stage_alist ~f:begin fun i (stage_name, ast_stage) ->
      let st_assigns : Vlog.assign list =
        let complete_ast_ep ep =
          { ep with ep_path = stage_name :: ep.ep_path }
        in
        let unlabeled_ast_wires, labeled_ast_wires =
          List.partition_map ast_stage.stage_wires
            ~f:begin fun (src_ast_drv, dst_ast_ep, ast_labels_opt) ->
              let src_ast_drv' =
                match src_ast_drv with
                | Ast_driver_var src_ast_ep ->
                    Ast_driver_var (complete_ast_ep src_ast_ep)
                | Ast_driver_const _ -> src_ast_drv
              in
              let dst_ast_ep' = complete_ast_ep dst_ast_ep in
              match ast_labels_opt with
              | None -> `Fst (src_ast_drv', dst_ast_ep')
              | Some ast_labels -> `Snd (src_ast_drv', dst_ast_ep', ast_labels)
            end
        in
        let assigns_unlabeled =
          List.map unlabeled_ast_wires ~f:(convert_wire top_node)
        in
        let assigns_labeled =
          let labeled_wires =
            List.map labeled_ast_wires
              ~f:begin fun (src_ast_drv, dst_ast_ep, ast_labels) ->
                let src_drv, dst_ep, width =
                  resolve_wire top_node (src_ast_drv, dst_ast_ep)
                in
                let label_set = resolve_labels lookup_label ast_labels in
                src_drv, dst_ep, width, label_set
              end
          in
          let expanded_labeled_wires = expand_datapaths labeled_wires in
          List.fold expanded_labeled_wires
            ~init:Endpoint.Map.empty
            ~f:begin fun map (src_drv, dst_ep, label_set) ->
              Endpoint.Map.update map dst_ep ~f:begin fun sources_opt ->
                match sources_opt with
                | Some sources -> (src_drv, label_set) :: sources
                | None -> [src_drv, label_set]
              end
            end
          |> Endpoint.Map.to_alist
          |> List.map ~f:begin fun (dst_ep, sources) ->
              let dst_path, dst_bit = dst_ep in
              let assign_lhs = string_of_path dst_path, (dst_bit, dst_bit+1) in
              let assign_rhs =
                List.map sources ~f:begin fun (src_drv, label_set) ->
                  let label_mask_expr =
                    List.map (Int.Set.to_list label_set) ~f:begin fun i ->
                      Vlog.Sym_part begin
                        string_of_path [stage_name;"inst_class"], (i, i+1)
                      end
                    end
                    |> Vlog.make_binary_expr Vlog.Or
                  in
                  Vlog.Binary begin
                    Vlog.And,
                    label_mask_expr,
                    begin match src_drv with
                    | Driver_const src_value -> Vlog.Literal (1, src_value)
                    | Driver_var src_ep -> vlog_of_endpoint src_ep
                    end
                  end
                end
                |> Vlog.make_binary_expr Vlog.Or
              in
              assign_lhs, assign_rhs
          end
        in
        assigns_labeled @ assigns_unlabeled
      in
      let stage_node = stage_node_arr.(i) in
      let st_instances : (string * Vlog.vlog_instance) list =
        List.map ast_stage.stage_modules ~f:begin fun (instance_name, (class_name, _)) ->
          let bbox_node =
            match stage_node.lookup instance_name with
            | Submodule node -> node
            | _ -> assert false
          in
          let ports =
            bbox_node.component_array
            |> Array.to_list
            |> List.map ~f:begin fun (comp_name, comp) ->
                match comp with
                | Port port ->
                    let port_name = comp_name in
                    port_name, port, Vlog.Sym (string_of_path [stage_name;instance_name;port_name])
                | _ -> assert false
            end
          in
          let vlog_inst = { Vlog.class_name; parameters = []; ports } in
          string_of_path [stage_name;instance_name], vlog_inst
        end
      in
      let st_wires : (string * int) list =
        map_ports (fun (path, port) -> string_of_path (stage_name :: path), fst port) stage_node
      in
      { st_wires; st_instances; st_assigns }
    end
  in
  let (top_instances : (string * Vlog.vlog_instance) list),
      (top_wires_tmp : (string * int) list list) =
    List.map top_blackboxes ~f:begin fun (bbox_name, bbox) ->
      let ports, wires =
        List.map bbox.ports ~f:begin fun (port_name, port) ->
          let wire_name = string_of_path [bbox_name;port_name] in
          (port_name, port, Vlog.Sym wire_name),
          (wire_name, fst port)
        end
        |> List.unzip
      in
      let ports =
        if bbox.has_clock
        then ("clk", (1, Vlog.Input), Vlog.Sym "clk") :: ports
        else ports
      in
      let vlog_inst =
        { Vlog.class_name = bbox.class_name;
          parameters = [];
          ports }
      in
      (bbox_name, vlog_inst), wires
    end
    |> List.unzip
  in
  let top_wires = List.concat top_wires_tmp in
  let wire_alist =
    List.concat (top_wires :: List.map stages ~f:(fun stage -> stage.st_wires))
  in
  let instance_alist =
    List.concat (top_instances :: List.map stages ~f:(fun stage -> stage.st_instances))
  in
  let reg_alist, dst_reg_assigns =
    let stage_has_fb : bool array =
      List.map sorted_ast_stage_alist ~f:begin fun (_, stage_ast) ->
        match stage_ast.stage_feedback_wires_opt with
        | None -> false
        | Some _ -> true
      end
      |> Array.of_list
    in
    let non_fb_datapaths, fb_datapaths =
      let resolve_datapaths =
        List.map ~f:(resolve_datapath top_node lookup_label)
      in
      let synthetic_ast_datapaths : ast_datapath list =
        List.map (List.tl_exn stage_names) ~f:begin fun stage_name ->
          let ast_src =
            { ep_path = [first_stage_name;"inst_class"];
              ep_range = None }
          in
          let ast_dst =
            { ep_path = [stage_name;"inst_class"];
              ep_range = None }
          in
          let ast_labels_opt = None in
          Ast_driver_var ast_src, ast_dst, ast_labels_opt
        end
      in
      let fb_ast_datapaths =
        List.map sorted_ast_stage_alist ~f:begin fun (stage_name, stage_ast) ->
          match stage_ast.stage_feedback_wires_opt with
          | None -> []
          | Some (fb_paths_ast : ast_datapath list) ->
              List.map fb_paths_ast ~f:begin fun fb_path_ast ->
                let complete_ast_ep ep =
                  { ep with ep_path = stage_name :: ep.ep_path }
                in
                let src_drv_ast, dst_ep_ast, labels_opt = fb_path_ast in
                let src_drv_ast' =
                  match src_drv_ast with
                  | Ast_driver_var src_ep_ast ->
                      let src_ep_ast' = complete_ast_ep src_ep_ast in
                      Ast_driver_var src_ep_ast'
                  | Ast_driver_const (src_width, src_value) ->
                      Ast_driver_const (src_width, src_value)
                in
                let dst_ep_ast' = complete_ast_ep dst_ep_ast in
                src_drv_ast', dst_ep_ast', labels_opt
              end
        end
        |> List.concat
      in
      resolve_datapaths (synthetic_ast_datapaths @ top_ast.top_paths),
      resolve_datapaths fb_ast_datapaths
    in
    let expanded_non_fb_datapaths = expand_datapaths non_fb_datapaths in
    let expanded_fb_datapaths = expand_datapaths fb_datapaths in
    let via_regs_of_datapath (dp_exp : datapath1) : (string * endpoint) list =
      let src_drv, dst_ep, _ = dp_exp in
      match src_drv with
      | Driver_const _ -> []
      | Driver_var src_ep ->
          let src_stage_name = List.hd_exn (fst src_ep) in
          let dst_stage_name = List.hd_exn (fst dst_ep) in
          let src_stage_index = index_of_stage src_stage_name in
          let dst_stage_index = index_of_stage dst_stage_name in
          List.range (src_stage_index+1) dst_stage_index
          |> List.map ~f:begin fun via_stage_index ->
            let via_stage_name = stage_name_arr.(via_stage_index) in
            via_stage_name, src_ep
          end
    in
    let we_expr_of stage_name =
      Vlog.Sym (string_of_path ["pipectl";sprintf "o_%s_we" stage_name]) (* TODO: hard-coded *)
    in
    let via_ir_name_of stage_name ep =
      let path, bit = ep in
      "VR" ^ string_of_endpoint (stage_name :: path, bit)
    in
    let fb_expr_of stage_name =
      Vlog.Sym (string_of_path [stage_name;"o_fb"])
    in
    let via_regs =
      List.map expanded_non_fb_datapaths ~f:via_regs_of_datapath
      |> List.concat
      |> List.dedup
      |> List.map ~f:begin fun (via_stage_name, src_ep) ->
        let src_path, src_bit = src_ep in
        let src_stage_name = List.hd_exn src_path in
        let via_stage_index = index_of_stage via_stage_name in
        let next_expr =
          let src_stage_index = index_of_stage src_stage_name in
          let prev_stage_index = via_stage_index-1 in
          if prev_stage_index = src_stage_index
          then
            Vlog.Sym_part (string_of_path src_path, (src_bit, src_bit+1))
          else
            let prev_stage_name = stage_name_arr.(prev_stage_index) in
            Vlog.Sym (via_ir_name_of prev_stage_name src_ep)
        in
        let we_expr_opt =
          if stage_has_fb.(via_stage_index)
          then
            let via_fb_expr = fb_expr_of via_stage_name in
            Some (Vlog.Binary begin
              Vlog.And,
              we_expr_of via_stage_name,
              Vlog.Unary (Vlog.Not, via_fb_expr)
            end)
          else
            Some (we_expr_of via_stage_name)
        in
        via_ir_name_of via_stage_name src_ep, (1, next_expr, we_expr_opt)
      end
    in
    let dst_regs, dst_reg_assigns =
      begin
        List.map expanded_non_fb_datapaths ~f:(fun dp -> (false, dp)) @
        List.map expanded_fb_datapaths ~f:(fun dp -> (true, dp))
      end
      |> List.fold
        ~init:Endpoint.Map.empty
        ~f:begin fun map (is_fb, (src_drv, dst_ep, label_opt)) ->
          Endpoint.Map.update map dst_ep ~f:begin fun sources_opt ->
            match sources_opt with
            | Some sources -> (is_fb, src_drv, label_opt) :: sources
            | None -> [is_fb, src_drv, label_opt]
          end
        end
      |> Endpoint.Map.to_alist
      |> List.map ~f:begin fun (dst_ep, sources_all) ->
        let dst_path, dst_bit = dst_ep in
        let dst_stage_name = List.hd_exn dst_path in
        let dst_stage_index = index_of_stage dst_stage_name in
        let sources_fb_opt, sources_non_fb =
          let sources_fb, sources_non_fb =
            List.partition_map sources_all ~f:begin fun (is_fb, src_drv, label_opt) ->
              if is_fb
              then `Fst (src_drv, label_opt)
              else `Snd (src_drv, label_opt)
            end
          in
          begin
            if stage_has_fb.(dst_stage_index)
            then Some sources_fb
            else None
          end,
          sources_non_fb
        in
        let dst_reg_name = "DR" ^ string_of_endpoint dst_ep in
        let dst_reg_assign =
          (string_of_path dst_path, (dst_bit, dst_bit+1)), Vlog.Sym dst_reg_name
        in
        let dst_fb_expr = fb_expr_of dst_stage_name in
        let dst_reg_we =
          match sources_fb_opt with
          | Some [] ->
            Some (Vlog.Binary begin
              Vlog.And,
              we_expr_of dst_stage_name,
              Vlog.Unary (Vlog.Not, dst_fb_expr)
            end)
          | _ ->
            Some (we_expr_of dst_stage_name)
        in
        let prev_stage_index = dst_stage_index-1 in
        let prev_stage_name = stage_name_arr.(prev_stage_index) in
        let expr_of_src_drv = function
          | Driver_const src_value -> Vlog.Literal (1, src_value)
          | Driver_var src_ep ->
              let src_path, src_bit = src_ep in
              let src_stage_name = List.hd_exn src_path in
              let src_stage_index = index_of_stage src_stage_name in
              if src_stage_index >= prev_stage_index
              then Vlog.Sym_part (string_of_path src_path, (src_bit, src_bit+1))
              else Vlog.Sym (via_ir_name_of prev_stage_name src_ep)
        in
        let dst_reg_next_of is_fb sources =
          begin match sources with
          | [src_drv, None] -> Some src_drv
          | _ -> None
          end
          |> function
          (* single unlabeled source *)
          | Some src_drv -> expr_of_src_drv src_drv
          (* potentially multiple labeled sources *)
          | None ->
              if sources = []
              then Vlog.Literal (1, 0)
              else begin
                List.map sources ~f:begin fun (src_drv, label_opt) ->
                  let label =
                    match label_opt with
                    | Some label -> label
                    | None ->
                        failwithf "datapath destination conflict on %s"
                          dst_reg_name ()
                  in
                  let label_mask_expr =
                    List.map (Int.Set.to_list label) ~f:begin fun i ->
                      let src_stage_name =
                        if is_fb
                        then dst_stage_name
                        else prev_stage_name
                      in
                      Vlog.Sym_part begin
                        string_of_path [src_stage_name;"inst_class"], (i, i+1)
                      end
                    end
                    |> Vlog.make_binary_expr Vlog.Or
                  in
                  Vlog.Binary begin
                    Vlog.And,
                    label_mask_expr,
                    expr_of_src_drv src_drv
                  end
                end
                |> Vlog.make_binary_expr Vlog.Or
              end
        in
        let dst_reg_next_non_fb = dst_reg_next_of false sources_non_fb in
        let dst_reg_next =
          match sources_fb_opt with
          | None | Some [] ->
              dst_reg_next_non_fb
          | Some sources_fb (* non-empty *) ->
              let dst_reg_next_fb = dst_reg_next_of true sources_fb in
              Vlog.Cond begin
                dst_fb_expr,
                dst_reg_next_fb,
                dst_reg_next_non_fb
              end
        in
        let dst_reg = dst_reg_name, (1, dst_reg_next, dst_reg_we) in
        dst_reg, dst_reg_assign
      end
      |> List.unzip
    in
    dst_regs @ via_regs, dst_reg_assigns
  in
  let assigns =
    let top_assigns = List.map top_ast.top_wires ~f:(convert_wire top_node) in
    List.concat (dst_reg_assigns :: top_assigns :: List.map stages ~f:(fun stage -> stage.st_assigns))
  in
  let top_vlog_module =
    { Vlog.port_alist = top_ast.top_ports;
      wire_alist;
      assigns;
      instance_alist;
      reg_alist;
      has_clock = true }
  in
  let vlog_module_name = "cpu" in
  Vlog.write_named_vlog_module
    (vlog_module_name, top_vlog_module)
    (sprintf "%s/%s.v" destdir vlog_module_name)

let () =
  Command.basic
    ~summary:"Generate Verilog code from SADL"
    Command.Spec.(
      empty
      +> flag "-o" (required string) ~doc:"DIR output directory"
      +> flag "-I" (listed string) ~doc:"DIR module search path"
    )
    main
  |> Command.run
